package main

import (
	"bytes"
	"compress/gzip"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"strings"

	log "github.com/sirupsen/logrus"

	"github.com/clbanning/mxj"
	"github.com/spf13/viper"
)

var replaceKeys = make(map[string]string)

//InitConf 初始化配置文件
func initConf(cfgFile string) {
	if _, err := os.Stat(cfgFile); os.IsNotExist(err) {
		log.Warnf("config file(%s) not exist", cfgFile)
	}
	viper.SetConfigType("toml")
	viper.SetConfigFile(cfgFile)
	err := viper.ReadInConfig()
	if err != nil {
		log.Warnf("read config file(%s) error, details: %s", viper.ConfigFileUsed(), err)
	}
	log.Infof("Loading config file: %v", cfgFile)
	//配置默认值，如果配置内容中没有指定，就使用以下值来作为配置值，给定默认值是一个让程序更健壮的办法
	viper.SetDefault("app.port", "8083")

	viper.SetDefault("proxy.host", "http://map.cityfun.com.cn")
	viper.SetDefault("proxy.port", "")

	viper.SetDefault("kvs.keys", []string{"ows:Get", "ServiceMetadataURL", "ResourceURL"})
	viper.SetDefault("kvs.subpaths", []string{"#attr.xlink:href.#text", "#attr.xlink:href.#text", "#attr.template.#text"})

	//init kvs
	keys := viper.GetStringSlice("kvs.keys")
	subpaths := viper.GetStringSlice("kvs.subpaths")
	if len(keys) != len(subpaths) {
		log.Fatal("kvs lenth error")
	}
	for i, k := range keys {
		replaceKeys[k] = subpaths[i]
	}
}

func main() {
	initConf("conf.toml")
	pHost := viper.GetString("proxy.host")
	pPort := viper.GetString("proxy.port")
	if !strings.HasPrefix(pHost, "http://") {
		pHost = "http://" + pHost
	}
	if pPort != "" {
		pPort = ":" + pPort
	}

	targetURL, err := url.Parse(pHost + pPort)
	if err != nil {
		log.Fatal(err)
	}
	proxy := NewSingleHostReverseProxy(targetURL)
	port := viper.GetString("app.port")
	log.Printf("wmts reverse proxy server serve at : 127.0.0.1:%s (*^_^*)", port)
	if err := http.ListenAndServe(":"+port, proxy); err != nil {
		log.Fatal("Start server failed,err:", err)
	}
}

//NewSingleHostReverseProxy xxx
func NewSingleHostReverseProxy(target *url.URL) *httputil.ReverseProxy {
	targetQuery := target.RawQuery
	director := func(req *http.Request) {
		req.URL.Scheme = target.Scheme
		req.URL.Host = target.Host
		req.URL.Path = singleJoiningSlash(target.Path, req.URL.Path)
		if targetQuery == "" || req.URL.RawQuery == "" {
			req.URL.RawQuery = targetQuery + req.URL.RawQuery
		} else {
			req.URL.RawQuery = targetQuery + "&" + req.URL.RawQuery
		}
		req.Header.Set("X-Forwarded-Host", req.Header.Get("Host"))
		req.Host = target.Host
	}

	// 自定义ModifyResponse
	modifyResp := func(resp *http.Response) error {
		//.xml结尾
		if strings.HasSuffix(resp.Request.URL.Path, ".xml") {
			log.Printf(resp.Request.URL.String())
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				return err
			}
			//需要先解压
			contentEncoding := resp.Header.Get("Content-Encoding")
			if "gzip" == contentEncoding {
				// decompress the encoded bytes
				r, err := gzip.NewReader(bytes.NewBuffer(body))
				if err != nil {
					return err
				}
				defer r.Close()
				body, err = ioutil.ReadAll(r)
				if err != nil {
					return err
				}
			}
			xr := bytes.NewReader(body)
			m, err := mxj.NewMapXmlSeqReader(xr)
			if mxj.NO_ROOT == err {
				m, err = mxj.NewMapXmlSeqReader(xr)
			}
			if err != nil {
				return nil
			}
			tHost := viper.GetString("replace.host")
			if strings.HasPrefix(tHost, "http://") {
				tHost = tHost[7:]
			}

			if "" == tHost {
				tHost = resp.Request.Host
			}

			for k, v := range replaceKeys {
				paths := m.PathsForKey(k)
				for _, path := range paths {
					err := ReplaceURL(&m, path, v, tHost)
					if err != nil {
						log.Printf("replace url for path(%s)'s attribute(%s) err: %s\n", path, v, err)
					}
				}
			}

			newBody, err := m.XmlSeqIndent("", "	")
			if err != nil {
				return err
			}
			var buf bytes.Buffer
			if "gzip" == contentEncoding {
				// compress the encoded bytes
				gw := gzip.NewWriter(&buf)
				defer gw.Close()
				_, err := gw.Write(newBody)
				if err != nil {
					log.Error(err)
				}
				// flush and close the writer
				if err = gw.Close(); err != nil {
					log.Error(err)
				}
				newBody = buf.Bytes()
			}
			// 修改返回内容及ContentLength
			resp.Body = ioutil.NopCloser(bytes.NewBuffer(newBody))
			resp.ContentLength = int64(len(newBody))
			resp.Header.Set("Content-Length", fmt.Sprint(len(newBody)))
			// resp.Header.Set("Content-Encoding", "identity")
		}
		//else do nothing
		return nil
	}
	// 传入自定义的ModifyResponse
	return &httputil.ReverseProxy{Director: director, ModifyResponse: modifyResp}
}

// ReplaceURL replace url host for the subpaths, eg. #attr.xlink:href.#text / #attr.template.#text
func ReplaceURL(m *mxj.Map, path string, subpath string, host string) error {
	// get the array of Items entries for the 'path'
	vals, err := m.ValuesForPath(path)
	if err != nil {
		return fmt.Errorf("ValuesForPath err: %s", err.Error())
	} else if len(vals) == 0 {
		return fmt.Errorf("no vals for path: %s", path)
	}
	// process each Items entry
	for _, v := range vals {
		vm, ok := v.(map[string]interface{})
		if !ok {
			return fmt.Errorf("assertion failed")
		}
		// extract the associated url
		urltext, _ := mxj.Map(vm).ValueForPathString(subpath)
		u, err := url.Parse(urltext)
		if err != nil {
			return fmt.Errorf("parse %s err: %s", subpath, err)
		}
		newurl := strings.Replace(urltext, u.Host, host, 1)
		if err = mxj.Map(vm).SetValueForPath(newurl, subpath); err != nil {
			return fmt.Errorf("SetValueFor %s err: %s", subpath, err)
		}
	}
	return nil
}

func singleJoiningSlash(a, b string) string {
	aslash := strings.HasSuffix(a, "/")
	bslash := strings.HasPrefix(b, "/")
	switch {
	case aslash && bslash:
		return a + b[1:]
	case !aslash && !bslash:
		return a + "/" + b
	}
	return a + b
}

//ReadFileXML 读文件修改
func ReadFileXML(pathfile string) {
	data, err := ioutil.ReadFile(pathfile)
	if err != nil {
		log.Fatalln(err)
	}
	m, err := mxj.NewMapXmlSeq(data, false)
	if err != nil {
		log.Fatalln(err)
	}
	keys := map[string]string{
		"ows:Get":            "#attr.xlink:href.#text",
		"ServiceMetadataURL": "#attr.xlink:href.#text",
		"ResourceURL":        "#attr.template.#text",
	}
	host := "map.cityfun.com.cn"
	for k, v := range keys {
		paths := m.PathsForKey(k)
		for _, path := range paths {
			err := ReplaceURL(&m, path, v, host)
			if err != nil {
				log.Printf("replace url for path(%s)'s attribute(%s) err: %s\n", path, v, err)
			}
		}
	}

	out, err := m.XmlSeqIndent("", "	")
	if err != nil {
		log.Fatalln(err)
	}
	err = ioutil.WriteFile("out.xml", out, os.ModePerm)
	if err != nil {
		log.Fatalln(err)
	}
}
